import isaacgym
import torch
from asyncio import FastChildWatcher
from code import interact
from collections import deque
import os
import time
from tkinter.tix import Tree
import numpy as np
from isaacgym import gymutil
from isaacgym import gymapi
from math import sqrt
import math
from sympy import false
import cv2
from draw import *
from pcgworker.PCGWorker import *

from wfc_vecenv_stable_baselines import PCGVecEnvStableBaselines

from stable_baselines3 import PPO
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.results_plotter import load_results, ts2xy
from stable_baselines3.common.monitor import Monitor
from stable_baselines3.common import results_plotter
from stable_baselines3.common.results_plotter import load_results, ts2xy, plot_results
from stable_baselines3.common.noise import NormalActionNoise
from stable_baselines3.common.callbacks import BaseCallback

import asyncio

from multiprocessing import Process, Queue, Pipe, Manager, Value, Array, Lock, Event, Pool
import threading


def test_isaac_sb3_proc(model_param, compute_device_id = 0, graphics_device_id = 0, cuda_device_id = 0,return_queue = None):
    
    os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3,4,5,6,7'

    print("before m_env")

    m_env = PCGVecEnvStableBaselines(headless_ = True, compute_device_id = compute_device_id, graphics_device_id = graphics_device_id)

    print("after m_env")

    model = PPO('CnnPolicy', env=m_env, batch_size = 1024, device = 'cuda:' + str(cuda_device_id))
    model.set_parameters(copy.deepcopy(model_param))

    # print(model_param)

    print("model parameters transfer complete")

    print("SB3 model on device : ", model.device)
    
    m_env.reset()
    all_actions = np.zeros(m_env.num_matrix_envs)
    for i in range(m_env.num_matrix_envs):
        all_actions[i] = -2

    for i in range(400):
        observation, reward, done, info = m_env.step(all_actions)
        # print(compute_device_id,i)

    model_cpu = PPO('CnnPolicy', env=m_env, batch_size = 1024, device = 'cpu:0')
    model_cpu.set_parameters(model.get_parameters())
    return_queue.put(copy.deepcopy(model_cpu.get_parameters()))

    print("device of PPO model_cpu : ", model_cpu.device)

    # release VecEnv
    m_env.close()

    print("exit test_isaac_sb3_proc ...")


def generate_boost_model(return_queue = None):

    os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3,4,5,6,7'

    m_env = PCGVecEnvStableBaselines(headless_ = True)

    # boost model
    model = PPO('CnnPolicy', env=m_env, batch_size = 1024, device = 'cuda:0')

    model_cpu = PPO('CnnPolicy', env=m_env, batch_size = 1024, device = 'cpu:0')
    model_cpu.set_parameters(model.get_parameters())
    return_queue.put(
        (copy.deepcopy(model_cpu.get_parameters()),
        123,
        456
        ))

    m_env.close()

    print("boost model training complete")


def save_model_proc(model_parameters, generation_id):

    os.environ['CUDA_VISIBLE_DEVICES']='0,1,2,3,4,5,6,7'
    LOGDIR = "./training_logs"

    m_env = PCGVecEnvStableBaselines(headless_ = True)
    m_env.reset()
    # save model to LOGDIR
    model = PPO('CnnPolicy', env=m_env, batch_size = 1024, device = 'cuda:0')
    model.set_parameters(copy.deepcopy(model_parameters))
    model.save(LOGDIR + "/model_" + str(generation_id))

    print("model saved to : ", LOGDIR + "/model_" + str(generation_id))

def save_model(model_parameters, generation_id):

    p = Process(target=save_model_proc, args=(model_parameters, generation_id))
    p.start()
    p.join()
    p.kill()


if __name__ == '__main__':

    print("------------------------------- Multi-processing test -------------------------------")

    print('Parent process %s.' % os.getpid())

    # TODO generate boost model here

    model_param = None
    model_param1 = None
    results = None

    boost_model_manager = multiprocessing.Manager()
    boost_model_queue = boost_model_manager.Queue()

    boost_model_process = Process(target=generate_boost_model, args=(boost_model_queue,))
    boost_model_process.start()
    boost_model_process.join()

    if boost_model_queue.empty() == False:
        print("boost model generate complete")
        results = copy.copy(boost_model_queue.get())
        boost_model_queue.task_done()
        print(results[1], results[2])
    
    boost_model_manager.shutdown()
    boost_model_process.kill()

    while True:

        manager = multiprocessing.Manager()
        return_q = manager.Queue()

        p0 = Process(target=test_isaac_sb3_proc, args=(results[0], 0,0,0,return_q,))
        p1 = Process(target=test_isaac_sb3_proc, args=(results[0], 3,3,3,return_q,))
        p2 = Process(target=test_isaac_sb3_proc, args=(results[0], 5,5,5,return_q,))
        p3 = Process(target=test_isaac_sb3_proc, args=(results[0], 7,7,7,return_q,))

        p0.start()
        p1.start()
        p2.start()
        p3.start()

        p0.join()
        p1.join()
        p2.join()
        p3.join()

        print("finish one loop")

        # get model parameters from queue
        for i in range(return_q.qsize()):
            results1 = copy.copy(return_q.get())
            return_q.task_done()
            # print(results[1], results[2])

        p0.kill()
        p1.kill()
        p2.kill()
        p3.kill()

        manager.shutdown()

        save_model(results1, 1)


        


